#include "ps.h"
#include <stdlib.h>
#include <float.h>
#include <math.h>
#include <time.h>
#include <stdio.h>
#include <string.h> // memset

#define DEBUG_PS_SHOW_ITERACTIONS
// #define DEBUG_PS_SHOW_PROBABILITIES
// #define DEBUG_PS_USED_PARTICLES
// #define DEBUG_PS_SHOW_NEW_PARTICLES


double
polar_box_muller_gaussian_generator(double mean, double variance)
{
	// Box-Muller gaussian generator
	// ref: http://en.wikipedia.org/wiki/Box_Muller_transform and
	// ref: http://www.design.caltech.edu/erik/Misc/Gaussian.html

	static int generate_step = 1;
	static double y1 = 0.0, y2 = 0.0;

	if (generate_step)
	{
		double x1, x2, w;

		do
		{
			x1 = 2.0 * ((double) rand() / ((double) RAND_MAX)) - 1.0;
			x2 = 2.0 * ((double) rand() / ((double) RAND_MAX)) - 1.0;
			w = x1 * x1 + x2 * x2;

		} while (w >= 1.0);

		w = sqrt((-2.0 * log(w)) / w);

		y1 = x1 * w;
		y2 = x2 * w;

		generate_step = 0;
		return (mean + variance * y1);
	}
	else
	{
		generate_step = 1;
		return (mean + variance * y2);
	}
}


ProbabilisticSearch::ProbabilisticSearch(double (*fitness_function)(double *particle, void *data), double **limits, int dim, void *data, int num_particles, int max_iteractions)
{
	srand(time(NULL));

	_fitness_function = fitness_function;
	_max_iteractions = max_iteractions;
	_num_particles = num_particles;
	_limits = limits;
	_data = data;
	_dim = dim;

	_alloc_stuff();
	_intialize_best_particle();
}


ProbabilisticSearch::~ProbabilisticSearch()
{
	_dealloc_stuff();
}


void
ProbabilisticSearch::Optimize()
{
	int num_iteractions = 0;

	_randomize_particles();

	while (num_iteractions < _max_iteractions)
	{
		_evaluate_fitness(_particles, _fitness);
		_update_best_particle();

		if (USE_GAUSSIAN_DISTRO)
			_update_particles_with_selective_gaussian_mixture();
		else // USE_RANDOM_DISTRO
			_randomize_particles();

		_error_evolution[num_iteractions] = _best_fitness;

#ifdef DEBUG_PS_SHOW_ITERACTIONS
		fprintf(stderr, "Iteraction: %d of %d GbestFitness: %lf\n",
			num_iteractions, _max_iteractions, -_best_fitness);
#endif

		num_iteractions++;
	}
}


double 
ProbabilisticSearch::GetBestFitness()
{
	return _best_fitness;
}


double*
ProbabilisticSearch::GetBestSolution()
{
	return _best_particle;
}


double*
ProbabilisticSearch::GetErrorEvolution()
{
	return _error_evolution;
}


void 
ProbabilisticSearch::_alloc_stuff()
{
	// TODO: organizar segundo a ordem que esta na definicao da classe
	_new_particles = (double **) calloc (_num_particles, sizeof(double*));
	_particles = (double **) calloc (_num_particles, sizeof(double*)); 

	_particle_std = (double *) calloc (_num_particles, sizeof(double));
	_particle_probability = (double *) calloc (_num_particles, sizeof(double));
	_fitness_of_new_particles = (double *) calloc (_num_particles, sizeof(double));
	_fitness = (double *) calloc (_num_particles, sizeof(double)); 
	_best_particle = (double *) calloc (_dim, sizeof(double)); 
	_error_evolution = (double *) calloc (_max_iteractions, sizeof(double)); 
	_selected_gaussians = (int *) calloc (_num_particles, sizeof(int));
	_gaussian_probability = (double *) calloc (_num_particles, sizeof(double));
	_gaussian_cumulative_probability = (double *) calloc (_num_particles, sizeof(double));

	for (int i = 0; i < _num_particles; i++)
	{
		_new_particles[i] = (double *) calloc (_dim, sizeof(double));
		_particles[i] = (double *) calloc (_dim, sizeof(double)); 
	}
}


void 
ProbabilisticSearch::_dealloc_stuff()
{
	// TODO: eu estou colocando varias variaveis novas, checar se estou dando free nelas
	// TODO: organizar segundo a ordem que esta na definicao da classe

	for (int i = 0; i < _num_particles; i++)
	{
		free(_new_particles[i]);
		free(_particles[i]);
	}

	free(_new_particles);
	free(_particles);

	free(_fitness);
	free(_best_particle);
	free(_particle_probability);
	free(_fitness_of_new_particles);
	free(_error_evolution);
	free(_particle_std);
	free(_selected_gaussians);
	free(_gaussian_probability);
	free(_gaussian_cumulative_probability);
}


void
ProbabilisticSearch::_randomize_particles()
{
	int i, j;
	double r; // random value between 0 and 1
	double ss_size; // search space size

	for (i = 0; i < _num_particles; i++)
	{
		for (j = 0; j < _dim; j++)
		{
			r = ((double) rand() / (double) RAND_MAX);
			ss_size = _limits[j][1] - _limits[j][0];
			_particles[i][j] = r * ss_size + _limits[j][0];

			if ((_particles[i][j] > _limits[j][1]) || (_particles[i][j] < _limits[j][0]))
				exit(printf("Error: the randomization is wrong...\n"));
		}
	}
}


void
ProbabilisticSearch::_intialize_best_particle()
{
	_best_fitness = -DBL_MAX;
	memset(_best_particle, 0, _dim * sizeof(double));
}


void
ProbabilisticSearch::_evaluate_fitness(double **particles, double *fitness)
{
	int i;

	// it's highly parallelizable!!
	for (i = 0; i < _num_particles; i++)
		fitness[i] = _fitness_function(particles[i], _data);
}


void
ProbabilisticSearch::_update_best_particle()
{
	int i;

	for (i = 0; i < _num_particles; i++)
	{
		if (_fitness[i] > _best_fitness)
		{
			_best_fitness = _fitness[i];
			_copy(_particles[i], _best_particle);			
		}
	}
}


double
ProbabilisticSearch::_get_min_fitness()
{
	double min_fitness = DBL_MAX;

	for (int i = 0; i < _num_particles; i++)
		if (_fitness[i] < min_fitness)
			min_fitness = _fitness[i];

	return min_fitness;
}


int
ProbabilisticSearch::_get_particle_with_min_fitness()
{
	int min_fitness_particle = -1;
	double min_fitness = DBL_MAX;

	for (int i = 0; i < _num_particles; i++)
	{
		if (_fitness[i] < min_fitness)
		{
			min_fitness = _fitness[i];
			min_fitness_particle = i;
		}
	}

	return min_fitness_particle;
}


void
ProbabilisticSearch::_compute_particle_probability_from_fitness(double *fitness, double *probabilities)
{
	int i;
	double min_fitness;
	double sum_of_probabilities;

	min_fitness = _get_min_fitness();
	sum_of_probabilities = 0;

	// make all fitness bigger than zero
	for (i = 0; i < _num_particles; i++)
	{
		// I added 1e-10 to avoid zero-probability particles
		probabilities[i] = fitness[i] + fabs(min_fitness) + 1e-5;
		sum_of_probabilities += probabilities[i];
	}

	// normalize to [0, 1] interval
	for (i = 0; i < _num_particles; i++)
	{
#ifdef DEBUG_PS_SHOW_PROBABILITIES
		printf("fitness[%d]: % 15lf\t probability[%d]: % 15lf\t norm probability[%d]: %.20lf\n",
			i, fitness[i], i, probabilities[i], i, probabilities[i] / sum_of_probabilities);

		getchar();
#endif

		probabilities[i] /= sum_of_probabilities;
	}
}


void
ProbabilisticSearch::_calculate_stds_from_probabilities()
{
	double sqrt_2pi = sqrt(2 * M_PI);

	for (int i = 0; i < _num_particles; i++)
		// the calculation can be directly calculated from the normal
		// distribution formula, assuming that the particle is equal the mean
		_particle_std[i] = 1.0 / (_particle_probability[i] * sqrt_2pi);
}


double
ProbabilisticSearch::_compute_probability_of_particle_using_distribution(double *particle, int distribution_id)
{
	double covariance;
	double probability;
	double distance;

	distance = 0.0;

	// Note: distribution's mean is the particle value
	for (int i = 0; i < _dim; i++)
		distance += pow(particle[i] - _particles[distribution_id][i], 2);

	// Note: As I'm using the same std to all dimensions, the parameter of the exponential
	// is just -1/2 * 1/(std^2) * ((euclidean distance from particle to distribution's mean) ^ 2).
	// Access http://en.wikipedia.org/wiki/Multivariate_normal_distribution to a review on
	// multivariate normal distributions
	covariance = pow(_particle_std[distribution_id], 2);
	probability = (1.0 / sqrt(pow(2.0 * M_PI, _dim) * covariance)) * exp((-1.0 / 2.0) * (1.0 / covariance) * distance);

	return probability;
}


int
ProbabilisticSearch::_find_distribution_with_maximum_probability(double *particle)
{
	double probability;
	double max_probability;
	int selected_particle;

	max_probability = -DBL_MAX;
	selected_particle = -1;

	// find the particle's distribution with maximum
	// probability in the particle[i] point
	for (int i = 0; i < _num_particles; i++)
	{
		probability = _compute_probability_of_particle_using_distribution(particle, i);

		if (probability > max_probability)
		{
			max_probability = probability;
			selected_particle = i;
		}
	}

	return selected_particle;
}


void
ProbabilisticSearch::_set_distribution_as_used(int distribution)
{
	// check if the distribution is already in the list
	for (int i = 0; i < _num_selected_gaussians; i++)
		if (_selected_gaussians[i] == distribution)
			return;

	// if it's not, add it!
	_selected_gaussians[_num_selected_gaussians] = distribution;
	_num_selected_gaussians++;
}


void
ProbabilisticSearch::_select_particles_to_compose_distribution()
{
	int selected_distribution;

	_num_selected_gaussians = 0;

	for (int i = 0; i < _num_particles; i++)
	{
		selected_distribution = _find_distribution_with_maximum_probability(_particles[i]);
		_set_distribution_as_used(selected_distribution);
	}
}


void
ProbabilisticSearch::_view_selectede_particles_to_compose_distribution()
{
	int gaussian_used;
	FunctionAndPointsChart fitness_chart(_dim, _limits);
	GaussianMixtureChart gm_chart(_dim, _limits);

	fitness_chart.AddFunction(_fitness_function, NULL);

	for (int i = 0; i < _num_particles; i++)
	{
		fitness_chart.AddPoint(_particles[i], _fitness[i]);

		gaussian_used = 0;

		for (int j = 0; j < _num_selected_gaussians; j++)
		{
			if (_selected_gaussians[j] == i)
			{
				gaussian_used = 1;
				break;
			}
		}

		gm_chart.AddGaussian(_particles[i], pow(_particle_std[i], 2), gaussian_used);
	}


	printf("Num used gaussians: %d\n", _num_selected_gaussians);
	_viewer.RemoveAllCharts();
	_viewer.AddChart(&gm_chart);
	_viewer.AddChart(&fitness_chart);
	_viewer.View();
}


void
ProbabilisticSearch::_calculate_probability_of_selected_distributions()
{
	double accumulated = 0.0;
	double sum_of_probabilities = 0.0;

	for (int i = 0; i < _num_selected_gaussians; i++)
		sum_of_probabilities += _particle_probability[_selected_gaussians[i]];

	// normalize the probability of each gaussian
	for (int i = 0; i < _num_selected_gaussians; i++)
	{
		_gaussian_probability[i] = (_particle_probability[_selected_gaussians[i]] / sum_of_probabilities);

		accumulated += _gaussian_probability[i];
		_gaussian_cumulative_probability[i] = accumulated;
	}

#ifdef DEBUG_PS_USED_PARTICLES
	printf("_num_selected_gaussians: %d (%.2lf percent)\n", _num_selected_gaussians, 100.0 * ((double) _num_selected_gaussians / (double) _num_particles));

	for (int i = 0; i < _num_selected_gaussians; i++)
	{
		printf("\t-> _selected_gaussian: % 4d \t fitness: % 15lf \t probability: %.10lf \t accumulated: %.5lf\n", _selected_gaussians[i],
			_fitness[_selected_gaussians[i]], _gaussian_probability[i], _gaussian_cumulative_probability[i]);
	}

	getchar();
#endif
}


void
ProbabilisticSearch::_build_selective_mixture_of_gaussians(double *probailities)
{
	_calculate_stds_from_probabilities();
	_select_particles_to_compose_distribution();
	_view_selectede_particles_to_compose_distribution();
	_calculate_probability_of_selected_distributions();
}


int
ProbabilisticSearch::_sample_gaussian_from_mixture()
{
	// random between 0 and 1
	double r = ((double) rand()) / ((double) RAND_MAX);

#ifdef DEBUG_PS_USED_PARTICLES
	for (int i = 0; i < _num_selected_gaussians; i++)
		printf("_gaussian_cumulative_probability[%d]: %lf\n", i, _gaussian_cumulative_probability[i]);
#endif

	for (int i = 0; i < _num_selected_gaussians; i++)
	{
		if (r < _gaussian_cumulative_probability[i])
		{
#ifdef DEBUG_PS_USED_PARTICLES
			printf("r: %lf selected: %d\n\n", r, i);
			getchar();
#endif

			return i;
		}
	}

	exit(printf("Error::_sample_gaussian_from_mixture()::Invalid Gaussian\n"));
}


void
ProbabilisticSearch::_sample_particle_from_gaussian(int distribution_id, double *particle)
{
	double *mean;
	double variance;

	mean = _particles[_selected_gaussians[distribution_id]];
	variance = pow(1.0 / (_gaussian_probability[distribution_id] * sqrt(2 * M_PI)), 2);

#ifdef DEBUG_PS_SHOW_NEW_PARTICLES
	printf("selected_mean: %d variance: %lf\n", _selected_gaussians[distribution_id], variance);
#endif

	for (int i = 0; i < _dim; i++)
	{
		// sample a new particle value while the sampled value is out of limits
		do
		{
			particle[i] = polar_box_muller_gaussian_generator(mean[i], variance);

		} while ((particle[i] < _limits[i][0]) || (particle[i] > _limits[i][1]));
	}
}


void
ProbabilisticSearch::_sample_new_particles_from_distribution()
{
	int gaussian;

	for (int i = 0; i < _num_particles; i++)
	{
		gaussian = _sample_gaussian_from_mixture();
		_sample_particle_from_gaussian(gaussian, _new_particles[i]);
	}
}


void
ProbabilisticSearch::_replace_particles_with_better_fitness()
{
	int particle_to_replace;

	for (int i = 0; i < _num_particles; i++)
	{
		if (REPLACE_THE_WORST)
			particle_to_replace = _get_particle_with_min_fitness();
		else // check only the particle for improvements
			particle_to_replace = i;

		if (_fitness_of_new_particles[i] > _fitness[particle_to_replace])
		{
			// copy the new particle value to the particle to be replaced
			memcpy(_particles[particle_to_replace], _new_particles[i], _dim * sizeof(double));
			// copy the new particle fitness to the old particle fitness
			_fitness[particle_to_replace] = _fitness_of_new_particles[i];
		}
	}
}


void
ProbabilisticSearch::_update_particles_with_selective_gaussian_mixture()
{
	_compute_particle_probability_from_fitness(_fitness, _particle_probability);
	_build_selective_mixture_of_gaussians(_particle_probability);

#ifdef DEBUG_PS_SHOW_NEW_PARTICLES
	printf("\n\n");
#endif

	_sample_new_particles_from_distribution();

#ifdef DEBUG_PS_SHOW_NEW_PARTICLES
	printf("\n\n");
#endif

	_evaluate_fitness(_new_particles, _fitness_of_new_particles);

#ifdef DEBUG_PS_SHOW_NEW_PARTICLES
	printf("PARTICLES:\n");
	for (int i = 0; i < _num_particles; i++)
		printf("particle[%d]: % 10lf \t fitness: %lf\n", i, _particles[i][0], _fitness[i]);

	printf("\nNEW_PARTICLES:\n");
	for (int i = 0; i < _num_particles; i++)
		printf("particle[%d]: % 10lf \t fitness: %lf\n", i, _new_particles[i][0], _fitness_of_new_particles[i]);

	getchar();
#endif

	_replace_particles_with_better_fitness();
}


void
ProbabilisticSearch::_copy(double from[], double to[])
{
	int i;

	for (i = 0; i < _dim; i++)
	{
		to[i] = from[i];
	}
}

